package com.hrh.hrhaiagent.agent;

import cn.hutool.core.collection.CollUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.hrh.hrhaiagent.agent.model.AgentState;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.List;
import java.util.stream.Collectors;

@Data
@Slf4j
public class ToolCallAgent extends ReActAgent{
    // 工具调用
    private final ToolCallback[] availableTools;

    // 工具调用结果
    private ChatResponse toolCallChatResponse;

    // 工具调用管理器
    private final ToolCallingManager toolCallingManager;

    // 工具调用参数
    private final ChatOptions chatOptions;

    public ToolCallAgent(ToolCallback[] availableTools){
        super();
        this.availableTools = availableTools;
        this.toolCallingManager = ToolCallingManager.builder().build();
        this.chatOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true)
                .build();
    }

    /**
     * 思考
     * @return
     */
    @Override
    public boolean think() {
        if (getNextStepPrompt() != null && !getNextStepPrompt().isEmpty()){
            UserMessage userMessage = new UserMessage(getNextStepPrompt());
            addMessage(userMessage);
        }
        List<Message> messageList = getMessages();
        Prompt prompt = new Prompt(messageList, chatOptions);
        try{
            ChatResponse chatResponse = getChatClient().prompt(prompt)
                    .system(getSystemPrompt())
                    .tools(availableTools)
                    .call()
                    .chatResponse();
            this.toolCallChatResponse = chatResponse;
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
            String result = assistantMessage.getText();
            List<AssistantMessage.ToolCall> toolCallList = assistantMessage.getToolCalls();
            log.info("[{}] 思考内容: {}", getName(), result);
            log.info("[{}] 选择了 {} 个工具", getName(), toolCallList.size());
            if (toolCallList.isEmpty()){
                addMessage(assistantMessage);
                return false;
            } else {
                return true;
            }
        } catch (Exception e){
            log.error("[{}] think()发生异常: {}", getName(), e.getMessage(), e);
            addMessage(new AssistantMessage("思考发生错误: " + e.getMessage()));
            return false;
        }
    }

    /**
     * 执行工具调用并处理结果
     *
     * @return 执行结果
     */
    @Override
    public String act() {
        if (!toolCallChatResponse.hasToolCalls()) {
            return "没有工具调用";
        }
        // 调用工具
        Prompt prompt = new Prompt(getMessages(), chatOptions);
        try {
            ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallChatResponse);
            clearMessages();
            for (Message msg : toolExecutionResult.conversationHistory()) {
                addMessage(msg);
            }
            ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());
            String results = toolResponseMessage.getResponses().stream()
                    .map(response -> "工具 " + response.name() + " 完成任务，结果: " + response.responseData())
                    .collect(Collectors.joining("\n"));
            log.info("[{}] 工具调用结果: {}", getName(), results);
            return results;
        } catch (Exception e) {
            log.error("[{}] act()工具调用异常: {}", getName(), e.getMessage(), e);
            return "工具调用发生错误: " + e.getMessage();
        }
    }

}
